#pragma once

/*
 * allocator for std::vector and std::string.  
 */

#include <vector>
#include <cstddef>
#include <cstdlib>
#include <new>
#include <limits>
#include <mutex>
#include <string>
#include <array>

namespace anet {
    namespace utils {
        // memory pool.
        class CMemoryPool {
        private:
            struct Block {
                Block* next;
            };

            struct Chunk {
                void* memory;
                Block* blocks;
            };

            // align size.
            static constexpr size_t ALIGN = alignof(std::max_align_t);

            // max block size.
            static constexpr size_t MAX_BLOCK_SIZE = 1024;

            // chunk size.
            static constexpr size_t CHUNK_SIZE = 8192;

            // all count.
            static constexpr size_t ALL_COUNT = MAX_BLOCK_SIZE / ALIGN;

            std::vector<Chunk> m_chunks;
            Block* m_free_blocks[ALL_COUNT] = { nullptr };
            std::array<std::mutex, ALL_COUNT> m_mutexes;

        private:
            size_t align(size_t n) {
                return (n + ALIGN - 1) & ~(ALIGN - 1);
            }

            size_t index(size_t n) {
                return (n - 1) / ALIGN;
            }

        public:
            CMemoryPool() = default;
            ~CMemoryPool() {
                for (auto& chunk : m_chunks) {
                    std::free(chunk.memory);
                }
            }

            void* allocate(size_t size) {
                if (size == 0) return nullptr;
                if (size > MAX_BLOCK_SIZE) {
                    return std::malloc(size);
                }

                size = align(size);
                size_t Index = index(size);

                // lock guard.
                std::lock_guard<std::mutex> guard(m_mutexes[Index]);
                if (m_free_blocks[Index] == nullptr) {
                    try {
                        allocate_chunk(size);
                    } catch (...) {
                        return nullptr;
                    }
                }

                Block* block = m_free_blocks[Index];
                m_free_blocks[Index] = block->next;
                return block;
            }

            void deallocate(void* ptr, size_t size) {
                if (ptr == nullptr) {
                    return;
                }

                if (size > MAX_BLOCK_SIZE) {
                    std::free(ptr);
                    return;
                }

                size = align(size);
                size_t Index = index(size);

                Block* block = static_cast<Block*>(ptr);

                // lock guard.
                std::lock_guard<std::mutex> guard(m_mutexes[Index]);
                block->next = m_free_blocks[Index];
                m_free_blocks[Index] = block;
            }

        private:
            void allocate_chunk(size_t block_size) {
                Chunk chunk;
                chunk.memory = std::malloc(CHUNK_SIZE);
                if (chunk.memory == nullptr) {
                    throw std::bad_alloc();
                    return;
                }

                chunk.blocks = nullptr;
                size_t count = CHUNK_SIZE / block_size;
                char* start = static_cast<char*>(chunk.memory);
                for (size_t i = 0; i < count; ++i) {
                    Block* block = reinterpret_cast<Block*>(start + i * block_size);
                    block->next = chunk.blocks;
                    chunk.blocks = block;
                }

                m_free_blocks[index(block_size)] = chunk.blocks;
                m_chunks.push_back(chunk);
            }
        };

        // pool allocator.
		template <typename T>
		class CPoolAllocator {
		public:
			using value_type = T;
			using pointer = T*;
			using const_pointer = const T*;
			using reference = T&;
			using const_reference = const T&;
			using size_type = std::size_t;
			using difference_type = std::ptrdiff_t;

			template <class U>
			struct rebind {
				typedef CPoolAllocator<U> other;
			};

			CPoolAllocator() noexcept {}
			template <class U> CPoolAllocator(const CPoolAllocator<U>&) noexcept {}

			pointer allocate(std::size_t n) {
				return static_cast<pointer>(get_memory_pool().allocate(n * sizeof(T)));
			}

			void deallocate(pointer p, std::size_t n) noexcept {
				get_memory_pool().deallocate(p, n * sizeof(T));
			}

			size_type max_size() const noexcept {
				return static_cast<size_type>(-1) / sizeof(value_type);
			}

			template<typename U, typename... Args>
			void construct(U* p, Args&&... args) {
				::new((void*)p) U(std::forward<Args>(args)...);
			}

			template<typename U>
			void destroy(U* p) {
				p->~U();
			}

			bool operator==(const CPoolAllocator&) const noexcept {
				return true;  
			}
			bool operator!=(const CPoolAllocator&) const noexcept {
				return false; 
			}

			static CMemoryPool& get_memory_pool() {
				static CMemoryPool pool;
				return pool;
			}
		};

        // operator == and != definition.
		template <typename T, typename U>
		bool operator==(const CPoolAllocator<T>&, const CPoolAllocator<U>&) noexcept {
			return true;  
		}
		template <typename T, typename U>
		bool operator!=(const CPoolAllocator<T>&, const CPoolAllocator<U>&) noexcept {
			return false;  
		}

		// define a new vector type with CPoolAllocator allocator.
		template <typename T>
		using PoolVector = std::vector<T, CPoolAllocator<T>>;

		// define a new string type with CPoolAllocator allocator.
		using PoolString = std::basic_string<char, std::char_traits<char>, CPoolAllocator<char>>;
    }
}
